"""
    https://console.bce.baidu.com/iam/#/iam/apikey/list
    API key type is 千帆ModelBuilder
    示例代码：https://cloud.baidu.com/doc/qianfan-api/s/3m7of64lb
"""
import requests
import json

import os

    
class QianFanModel:
    def __init__(self, model='ernie-4.5-turbo-32k'):
        self.model = model
        self.url = "https://qianfan.baidubce.com/v2/chat/completions"
        self.api_key = os.environ.get("QIANFAN_API_KEY")

    def chat(self, prompt):
        messages = [{'role': 'user', 'content': prompt}]
        
        payload = json.dumps({
            "model": self.model,
            "messages": messages
        })
        
        headers = {
            'Content-Type': 'application/json',
            'Authorization': 'Bearer ' + self.api_key,
        }
        
        response = requests.request("POST", self.url, headers=headers, data=payload)
        
        return response.json()['choices'][0]['message']['content']

if __name__ == '__main__':
    llm = QianFanModel()
    prompt_prefix = '按行判断下面每行文字是什么语言。然后将它翻译成英语和简体中文。最后按csv格式输出结果：第一列为原文，第二列为2字母的语言代码，第三列为英语结果，第四列为中文结果。如果原文或译文包含逗号或者引号，按照csv的标准方式进行转义处理。结果只需要输出一遍。不要输出csv表头。如果文字无法翻译，就将这一行舍弃。'
    prompt_content = """
una mente libera
uncinetto che conta
kit gilet jaune et triangle de signalisation voiture
cavo antifurto casco
tete de débroussailleuse husqvarna
portachiavi juventus ufficiale 2024
manette ps5 grise
2 bedroom tiny house
personalised stanley cup
"""
    result = llm.chat(prompt_prefix + '\n' + prompt_content)
    print(result)
